import pod5 as p5
import os
import argparse
from collections import defaultdict

def split_pod5_file(input_file, output_dir, read_count_dict):
    """拆分单个POD5文件到指定目录"""
    # 打开POD5文件
    with p5.Reader(input_file) as reader:
        # 遍历每条read
        for read in reader.reads():
            # 获取read_id并更新计数器
            read_id = read.read_id
            read_count_dict[read_id] += 1
            # 生成文件名：read_id + _x，其中x是该read_id的序号
            output_file = os.path.join(output_dir, f"{read_id}_{read_count_dict[read_id]}.pod5")
            # 创建一个新的POD5文件并写入单条read
            with p5.Writer(output_file) as writer:
                r = p5.Read(
                    read_id=read.read_id,
                    read_number=read.read_number,
                    start_sample=read.start_sample,
                    median_before=read.median_before,
                    end_reason=read.end_reason,
                    calibration=read.calibration,
                    pore=read.pore,
                    signal=read.signal,
                    run_info=read.run_info
                )
                writer.add_read(r)

    print(f"完成拆分：{input_file}")

def split_pod5_directory(input_dir, output_dir):
    """处理输入文件夹中的所有POD5文件"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 使用字典记录每个read_id的出现次数
    read_count_dict = defaultdict(int)

    # 遍历输入目录中的所有文件
    for filename in os.listdir(input_dir):
        if filename.endswith(".pod5"):  # 只处理POD5文件
            input_file = os.path.join(input_dir, filename)
            split_pod5_file(input_file, output_dir, read_count_dict)

    print(f"所有文件拆分完成！输出目录：{output_dir}")

def main():
    # 设置命令行参数解析器
    parser = argparse.ArgumentParser(description="将POD5文件夹中的文件拆分成一个文件一条read")
    parser.add_argument("-i", "--input", required=True, help="输入POD5文件夹路径")
    parser.add_argument("-o", "--output", default="split_reads", help="输出文件夹路径（默认：split_reads）")

    # 解析参数
    args = parser.parse_args()

    # 检查输入是否为目录
    if not os.path.isdir(args.input):
        print(f"错误：{args.input} 不是一个文件夹")
        return

    # 调用拆分函数
    split_pod5_directory(args.input, args.output)

if __name__ == "__main__":
    main()